# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question2/models.py
import torch.nn


# noinspection SpellCheckingInspection
class MultipoleDispersionModelTerm(torch.nn.Module):
    def __init__(self, random: torch.Generator):
        super().__init__()
        self.gamma0 = torch.nn.Parameter(torch.rand([], generator=random))
        self.gamma1 = torch.nn.Parameter(torch.rand([], generator=random))
        self.delta0 = torch.nn.Parameter(torch.rand([], generator=random))
        self.delta1 = torch.nn.Parameter(torch.rand([], generator=random))
        self.__j = torch.tensor(1j)
        self.__two = torch.tensor(2.)

    def forward(self, omega):
        numerator = self.gamma0 + self.__j * omega * self.gamma1
        denominator = self.delta0 + self.__j * omega * self.delta1 - omega ** self.__two
        return numerator / denominator


# noinspection SpellCheckingInspection
class MultipoleDispersionModel(torch.nn.Module):
    # 此默认值是由 main_guess.py 给出的
    def __init__(self, term_count: int = 6, random: torch.Generator = None):
        super().__init__()
        if random is None:
            random = torch.Generator()
            random.manual_seed(1234)
        self.epsilon_infinity = torch.nn.Parameter(torch.tensor(10.))
        self.terms = torch.nn.ModuleList(
            (MultipoleDispersionModelTerm(random) for _ in range(term_count)))

    def forward(self, omega):
        result = self.epsilon_infinity
        for term in self.terms:
            result = result + term(omega)
        return torch.stack([torch.real(result), torch.imag(result)], dim=-1)
